#ifndef MATRIX_H
#define MATRIX_H

#include <vector>
#include <stdexcept>
#include <cmath>
#include <cstddef>  // for size_t

// Forward declarations
template<size_t Rows, size_t Cols> class Matrix;

// Vector type aliases
using Vector2 = Matrix<2, 1>;    // 2D position vector
using Vector4 = Matrix<4, 1>;    // State vector [x, y, vx, vy]

// Matrix type aliases
using Matrix2x2 = Matrix<2, 2>;  // 2D square matrix
using Matrix4x4 = Matrix<4, 4>;  // 4D square matrix
using Matrix2x4 = Matrix<2, 4>;  // Measurement matrix
using Matrix4x2 = Matrix<4, 2>;  // Kalman gain matrix

template<size_t Rows, size_t Cols>
class Matrix {
private:
    std::vector<double> data;

public:
    // Default constructor (zero matrix)
    Matrix() : data(Rows * Cols, 0.0) {}

    // Constructor with initial value
    explicit Matrix(double value) : data(Rows * Cols, value) {}

    // Constructor with vector (for Vector types)
    explicit Matrix(const std::vector<double>& values) {
        if (values.size() != Rows * Cols) {
            throw std::runtime_error("Invalid vector size for matrix construction");
        }
        data = values;
    }

    // Copy constructor
    Matrix(const Matrix& other) : data(other.data) {}

    // Access operators
    double& operator()(size_t i, size_t j) { return data[i * Cols + j]; }
    const double& operator()(size_t i, size_t j) const { return data[i * Cols + j]; }

    // For vector types, allow single index access
    template<size_t R = Rows, size_t C = Cols>
    typename std::enable_if<C == 1, double&>::type
    operator()(size_t i) { return data[i]; }

    template<size_t R = Rows, size_t C = Cols>
    typename std::enable_if<C == 1, const double&>::type
    operator()(size_t i) const { return data[i]; }

    // Matrix operations
    Matrix operator+(const Matrix& other) const {
        Matrix result;
        for (size_t i = 0; i < Rows * Cols; ++i) {
            result.data[i] = data[i] + other.data[i];
        }
        return result;
    }

    Matrix operator-(const Matrix& other) const {
        Matrix result;
        for (size_t i = 0; i < Rows * Cols; ++i) {
            result.data[i] = data[i] - other.data[i];
        }
        return result;
    }

    Matrix operator*(double scalar) const {
        Matrix result;
        for (size_t i = 0; i < Rows * Cols; ++i) {
            result.data[i] = data[i] * scalar;
        }
        return result;
    }

    template<size_t OtherCols>
    Matrix<Rows, OtherCols> operator*(const Matrix<Cols, OtherCols>& other) const {
        Matrix<Rows, OtherCols> result;
        for (size_t i = 0; i < Rows; ++i) {
            for (size_t j = 0; j < OtherCols; ++j) {
                double sum = 0.0;
                for (size_t k = 0; k < Cols; ++k) {
                    sum += (*this)(i, k) * other(k, j);
                }
                result(i, j) = sum;
            }
        }
        return result;
    }

    Matrix<Cols, Rows> transpose() const {
        Matrix<Cols, Rows> result;
        for (size_t i = 0; i < Rows; ++i) {
            for (size_t j = 0; j < Cols; ++j) {
                result(j, i) = (*this)(i, j);
            }
        }
        return result;
    }

    static Matrix Zero() { return Matrix(); }

    static Matrix Identity() {
        static_assert(Rows == Cols, "Identity matrix must be square");
        Matrix result;
        for (size_t i = 0; i < Rows; ++i) {
            result(i, i) = 1.0;
        }
        return result;
    }

    Matrix inverse() const {
        static_assert(Rows == Cols, "Matrix must be square for inverse");
        
        std::vector<std::vector<double>> aug(Rows, std::vector<double>(2 * Cols, 0.0));
        
        for (size_t i = 0; i < Rows; ++i) {
            for (size_t j = 0; j < Cols; ++j) {
                aug[i][j] = (*this)(i, j);
            }
            aug[i][i + Cols] = 1.0;
        }

        for (size_t i = 0; i < Rows; ++i) {
            double pivot = aug[i][i];
            if (std::abs(pivot) < 1e-10) {
                throw std::runtime_error("Matrix is singular");
            }

            for (size_t j = 0; j < 2 * Cols; ++j) {
                aug[i][j] /= pivot;
            }

            for (size_t j = 0; j < Rows; ++j) {
                if (i != j) {
                    double factor = aug[j][i];
                    for (size_t k = 0; k < 2 * Cols; ++k) {
                        aug[j][k] -= factor * aug[i][k];
                    }
                }
            }
        }

        Matrix result;
        for (size_t i = 0; i < Rows; ++i) {
            for (size_t j = 0; j < Cols; ++j) {
                result(i, j) = aug[i][j + Cols];
            }
        }
        return result;
    }

    Matrix& operator=(const Matrix& other) {
        if (this != &other) {
            data = other.data;
        }
        return *this;
    }

    const std::vector<double>& getData() const { return data; }
};

// Non-member operator
template<size_t Rows, size_t Cols>
Matrix<Rows, Cols> operator*(double scalar, const Matrix<Rows, Cols>& matrix) {
    return matrix * scalar;
}

#endif // MATRIX_H 